import aiomysql
import os

async def new_pool(host='localhost', port=3306, user='root', password='', db=''):
    """
    创建aiomysql连接池
    :param host: 数据库主机地址
    :param port: 数据库端口
    :param user: 数据库用户名
    :param password: 数据库密码
    :param db: 数据库名称
    :return: aiomysql连接池对象
    """
    # 先创建临时连接检查数据库是否存在
    temp_pool = await aiomysql.create_pool(
        host=host,
        port=port,
        user=user,
        password=password
    )
    
    async with temp_pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(f"SHOW DATABASES LIKE '{db}'") if db else await cur.execute("SHOW DATABASES")
            result = await cur.fetchone()
            if not result:
                await cur.execute(f"CREATE DATABASE IF NOT EXISTS {db}")
    
    temp_pool.close()
    await temp_pool.wait_closed()
    
    # 创建正式连接池
    return await aiomysql.create_pool(
        host=host,
        port=port,
        user=user,
        password=password,
        db=db
    )

async def get_pool_config(pool:aiomysql.Pool) -> dict:
    """
    获取连接池的配置信息
    :param pool: aiomysql连接池对象
    :return: 包含host、port、user、db等连接参数的字典
    """
    return {
        'host': pool._conn_kwargs.get('host'),
        'port': pool._conn_kwargs.get('port'),
        'user': pool._conn_kwargs.get('user'),
        'password': pool._conn_kwargs.get('password'),
        'db': pool._conn_kwargs.get('db')
    }

async def new_pool_env():
    """
    从环境变量获取MySQL连接参数并创建连接池
    :return: aiomysql连接池对象
    """
    # 从环境变量获取MySQL连接参数
    from dotenv import load_dotenv
    load_dotenv()

    # 创建连接池
    return await new_pool(
        host=os.getenv('MYSQL_HOST', 'localhost'),
        port=int(os.getenv('MYSQL_PORT', '3306')),
        user=os.getenv('MYSQL_USER', 'root'),
        password=os.getenv('MYSQL_PASSWORD', ''),
        db=os.getenv('MYSQL_DB', 'test')
    )